'''
Author: Yan Jianhang
Date: 2022-04-12 19:51:21
LastEditors: Yan Jianhang
LastEditTime: 2022-04-29 00:01:09
FilePath: /STransGrasp/lib/tf/vis_points.py
Description: 可视化 grasp transfer 后的抓取点。(grasp_data/*/grasp_points*.npz)
'''

import pickle, argparse, os, yaml, sys
import numpy as np
import trimesh
sys.path.append('./')

parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, default='datasets')
parser.add_argument('--category', type=str, default='mug')
parser.add_argument('--mode', type=str, default='eval')
parser.add_argument('--inst', type=str, default='e6dedae946ff5265a95fb60c110b25aa')
parser.add_argument('--part', type=str, default='0')
parser.add_argument('--grasp_id', type=str, default='0')
opt = parser.parse_args()


if __name__ == "__main__":
    vis_file_name = f"grasp_{opt.grasp_id}_points_on_surface"  # "grasp_points": free points, "grasp_points_on_surface": surface points
    for file in os.listdir("grasp_data/{}/{}".format(opt.category, opt.mode)):
        opt.inst = file
        for part in os.listdir("grasp_data/{}/{}/{}".format(opt.category, opt.mode, file)):
            opt.part = part
            
            g_pts_file = "grasp_data/{}/{}/{}/{}/{}.npz".format(opt.category, opt.mode, opt.inst, opt.part, vis_file_name)
            obj_file   = "datasets/obj/{}/{}/{}/{}.ply".format(opt.category, opt.mode, opt.inst, opt.part)
            recon_file = f"dif/eval/{opt.category}/{opt.inst}/0/checkpoints/test.ply"
            
            g_pts = np.load(g_pts_file)
            mesh = trimesh.load_mesh(obj_file)
            mesh.apply_scale(g_pts["obj_scale"])
            mesh.visual.face_colors = np.array([0xAA, 0xAA, 0xAA, 200], dtype=np.uint8)
            recon_mesh = trimesh.load_mesh(recon_file)
            recon_mesh.apply_scale(g_pts["obj_scale"] * 0.5)
            recon_mesh.visual.face_colors = np.array([0x00, 0x00, 0xAA, 200], dtype=np.uint8)

            grasp_points = trimesh.PointCloud(g_pts["grasp_points"]*g_pts["obj_scale"]*0.5,
                                              colors=[0xFF, 0, 0, 255])
            
            trimesh.Scene([mesh, recon_mesh, grasp_points]).show()